#include <iostream>

namespace shr
{
    template<class T>
    struct avlTreeNode
    {
        T _data;
        avlTreeNode* _parent;
        avlTreeNode* _leftChild;
        avlTreeNode* _rightChild;
        int _bf;

        avlTreeNode(T data = T())
            :_data(data), _bf(0),
            _parent(nullptr),
            _leftChild(nullptr),
            _rightChild(nullptr)
        { }
    };

    template<class T>
    class avl_tree
    {
    private:
        typedef avlTreeNode<T> Node;

    public:
        avl_tree()
            :_root(nullptr)
        { }

        bool insert(const T& data)
        {
            Node* newNode = new Node(data);
            if(_root == nullptr) { _root = newNode; return true; }
            Node* curNode = _root;
            Node* parent = nullptr;
            while(curNode)
            {
                if(data < curNode->_data) {
                    parent = curNode;
                    curNode = curNode->_leftChild;
                }
                else if(data > curNode->_data) {
                    parent = curNode;
                    curNode = curNode->_rightChild;
                }
                else { return false; }
            }
            if(data < parent->_data) {
                parent->_leftChild = newNode;
                newNode->_parent = parent;
            }
            else // if(data > parent->_data) 
            {
                parent->_rightChild = newNode;
                newNode->_parent = parent;
            }

            curNode = newNode;
            while(parent)
            {
                if(curNode == parent->_leftChild) {
                    parent->_bf--;
                }
                else if(curNode == parent->_rightChild) {
                    parent->_bf++;
                }
                else { perror("insert err"); exit(1); }

                if(parent->_bf == 0) { 
                    break;
                }
                else if(parent->_bf == 1 || parent->_bf == -1) {
                    curNode = parent;
                    parent = parent->_parent;
                }
                else if(parent->_bf == 2 || parent->_bf == -2)
                {
                    //rotate
                    if(parent->_bf == 2 && curNode->_bf == 1) {
                        rotateLeft(parent);
                    }
                    else if(parent->_bf == -2 && curNode->_bf == -1) {
                        rotateRight(parent);
                    }
                    else if(parent->_bf == 2 && curNode->_bf == -1) {
                        rotateRightLeft(parent);
                    }
                    else if(parent->_bf == -2 && curNode->_bf == 1) {
                        rotateLeftRight(parent);
                    }
                    else { perror("insert err"); exit(1); }
                    break;
                }
                else { perror("insert err"); exit(1); }
            }
            return true;
        }

        void printTest()
        {
            _preorder_travel(_root);
        }

        ~avl_tree()
        {
            _destroy(_root);
        }

    private:
        void _preorder_travel(Node* root)
        {
            if(root == nullptr) { return; }
            std::cout << root->_data << ' ';
            _preorder_travel(root->_leftChild);
            _preorder_travel(root->_rightChild);
        }

        void rotateLeft(Node* parent)
        {
            //std::cout << "rotateLeft, val:" << parent->_data  << std::endl;
            Node* curNode = parent->_rightChild;
            Node* curLeft = curNode->_leftChild;
            
            Node* ppnode = parent->_parent;
            
            parent->_rightChild = curLeft;
            if(curLeft != nullptr) {
                curLeft->_parent = parent;
            }
            curNode->_leftChild = parent;
            parent->_parent = curNode;
            curNode->_bf = parent->_bf = 0; //调整cur和parent的平衡因子

            if(ppnode == nullptr) {
                _root = curNode;
            }
            else
            {
                if(parent == ppnode->_leftChild) {
                    ppnode->_leftChild = curNode;
                }
                else if(parent == ppnode->_rightChild) {
                    ppnode->_rightChild = curNode;
                }
                else { perror("insert err"); exit(1); }
            }
            curNode->_parent = ppnode;
        }

        void rotateRight(Node* parent)
        {
            //std::cout << "rotateRight" << std::endl;
            Node* curNode = parent->_leftChild;
            Node* curRight = curNode->_rightChild;
            Node* ppNode = parent->_parent;

            parent->_leftChild = curRight;
            if(curRight != nullptr) {
                curRight->_parent = parent;
            }
            curNode->_rightChild = parent;
            parent->_parent = curNode;
            curNode->_bf = parent->_bf = 0;

            if(ppNode == nullptr) {
                _root = curNode;
            }
            else
            {
                if(parent == ppNode->_leftChild) {
                    ppNode->_leftChild = curNode;
                }
                else if(parent == ppNode->_rightChild) {
                    ppNode->_rightChild = curNode;
                }
                else { perror("insert err"); exit(1); }
            }
            curNode->_parent = ppNode;
        }

        void rotateLeftRight(Node* parent)
        {
            //std::cout << "rotateLeftRight, val:" << parent->_data << std::endl;
            Node* curNode = parent->_leftChild;
            Node* curRight = curNode->_rightChild;
            int bf = curRight->_bf;

            rotateLeft(curNode);
            rotateRight(parent);
            if(bf == -1)
            {
                curRight->_bf = 0;
                curNode->_bf = 0;
                parent->_bf = 1;
            }
            else if(bf == 1)
            {
                curRight->_bf = 0;
                curNode->_bf = -1;
                parent->_bf = 0;
            }
            else if(bf == 0)
            {
                curRight->_bf = 0;
                curNode->_bf = 0;
                parent->_bf = 0;
            }
            else { perror("insert err"); exit(1); }
        }

        void rotateRightLeft(Node* parent)
        {
            //std::cout << "rotateRightLeft, val:" << parent->_data << std::endl;
        }

        void _destroy(Node* root)
        {
            if(root == nullptr) { return; }
            _destroy(root->_leftChild);
            _destroy(root->_rightChild);
            delete root;
        }

    private:
        Node* _root;
    };

}